
#include "gsl_ext_multimin.h"
#include "linmin.h"
#include "quasinewton_matrix.h"
#include "quasinewton_state.h"

#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>

static int bfgs_alloc(void *vstate, size_t n)
{
  bfgs_state_t *state = (bfgs_state_t *) vstate;

  state->H = gsl_matrix_alloc(n, n);
  state->tmpH = gsl_matrix_alloc(n, n);
  state->Bp = gsl_vector_alloc(n);
  state->Sq = gsl_vector_alloc(n);

  state->d = gsl_vector_alloc(n);
  state->p = gsl_vector_alloc(n);
  state->q = gsl_vector_alloc(n);

  state->x_plus = gsl_vector_alloc(n);
  state->gradient_plus = gsl_vector_alloc(n);

  return GSL_SUCCESS;
}

static int lbfgs_alloc(void *vstate, size_t n)
{
  bfgs_alloc(vstate, n);
  
  bfgs_state_t *state = (bfgs_state_t *) vstate;
  
  state->alpha_coeff = gsl_vector_alloc(5);
  state->P = gsl_matrix_alloc(5, n);
  state->Q = gsl_matrix_alloc(5, n);
}

static int bfgs_f_alloc(void *vstate, size_t n)
{
  bfgs_state_t *state = (bfgs_state_t *) vstate;
  
  state->line_minimizer = lineminimizer_alloc(lineminimizer_fletcher, n);
  bfgs_alloc(vstate, n);
  
  return GSL_SUCCESS;
}

static int bfgs_mt_alloc(void *vstate, size_t n)
{
  bfgs_state_t *state = (bfgs_state_t *) vstate;
  
  state->line_minimizer = lineminimizer_alloc(lineminimizer_morethuente, n);
  lineminimizer_morethuente_set_gamma(state->line_minimizer, 0.0);
  
  bfgs_alloc(vstate, n);
  
  return GSL_SUCCESS;
}

static int bfgs_set(void *vstate, gsl_multimin_function_fdf * fdf,
                    const gsl_vector * x, double *f, gsl_vector * gradient,
                    double step_size, double tol)
{
  bfgs_state_t *state = (bfgs_state_t *) vstate;
  
  GSL_MULTIMIN_FN_EVAL_F_DF (fdf, x, f, gradient);
  state->tol = tol;
  
  gsl_matrix_set_identity(state->H);
  
  return GSL_SUCCESS;
}

static void bfgs_free (void *vstate)
{
  bfgs_state_t *state = (bfgs_state_t *) vstate;
  
  gsl_vector_free(state->gradient_plus);
  gsl_vector_free(state->x_plus);
  gsl_vector_free(state->q);
  gsl_vector_free(state->p);
  gsl_vector_free(state->d);
  gsl_vector_free(state->Sq);
  gsl_vector_free(state->Bp);
  gsl_matrix_free(state->tmpH);
  gsl_matrix_free(state->H);
  lineminimizer_free(state->line_minimizer);
}

static int bfgs_hess_iterate(void *vstate, gsl_multimin_function_fdf *fdf,
                             gsl_vector *x, double *f,
                             gsl_vector *gradient, gsl_vector *dx)
{
  double D;
  
  bfgs_state_t *state = (bfgs_state_t *) vstate;
  
  /* Invert B by using Cholesky decomposition and solve */
  /* the quasi-Newton equation, */
  /* s = -inv(B).g */
  gsl_matrix_memcpy(state->tmpH, state->H);
    
  gsl_error_handler_t *old_handler = gsl_set_error_handler_off();
    
  if(gsl_linalg_cholesky_decomp(state->tmpH) != GSL_EDOM)
  {
    gsl_blas_dcopy(gradient, state->d);
    gsl_blas_dscal(-1.0, state->d);
    gsl_linalg_cholesky_svx(state->tmpH, state->d);
  }
  else
  {
    /* If the Cholesky decomposition fails, i.e. B is not positive */
    /* definite, replace it by identity matrix. Then s=-g. */
    gsl_matrix_set_identity(state->H);
    gsl_blas_dcopy(gradient, state->d);
    gsl_blas_dscal(-1.0, state->d);
  }
  
  gsl_set_error_handler(old_handler); 
  
  lineminimizer_minimize(state->line_minimizer, fdf, 1.0,
                         state->tol, *f, x, gradient,
                         state->d, &state->alpha,
                         state->x_plus, &state->f_plus, state->gradient_plus);
  
  /* p = x(k+1) - x(k) */
  gsl_blas_dcopy(state->x_plus, state->p);
  gsl_blas_daxpy(-1.0, x, state->p);
  
  /* q = g(k+1) - g(k) */
  gsl_blas_dcopy(state->gradient_plus, state->q);
  gsl_blas_daxpy(-1.0, gradient, state->q);
  
  gsl_blas_ddot(state->p, state->q, &D);
  if(D <= 0.0)
    /* If the curvature condition fails, B (or S) is set to identity. */
    gsl_matrix_set_identity(state->H);
  else
    /*sr1_hess_update(state->p, state->q, NULL, state->H);*/
    bfgs_hess_update(state->p, state->q, state->Bp, state->H);
  
  /* update the current iterate and its function and gradient values */
  gsl_blas_dcopy(state->x_plus, x);
  gsl_blas_dcopy(state->gradient_plus, gradient);
  *f = state->f_plus;
  gsl_blas_dcopy(state->p, dx);
  
  return GSL_SUCCESS;
}

static int bfgs_iterate(void *vstate, gsl_multimin_function_fdf * fdf,
                        gsl_vector * x, double *f,
                        gsl_vector * gradient, gsl_vector * dx)
{
  double D;
  
  bfgs_state_t *state = (bfgs_state_t *) vstate;
  
  /* compute search direction d=-Sg */
  gsl_blas_dsymv(CblasLower, -1.0, state->H, gradient, 0.0, state->d);
  
  lineminimizer_minimize(state->line_minimizer, fdf, 1.0,
                         state->tol, *f, x, gradient, state->d, &state->alpha,
                         state->x_plus, &state->f_plus, state->gradient_plus);
  
  /* p = x(k+1) - x(k) */
  gsl_blas_dcopy(state->x_plus, state->p);
  gsl_blas_daxpy(-1.0, x, state->p);
  
  /* q = g(k+1) - g(k) */
  gsl_blas_dcopy(state->gradient_plus, state->q);
  gsl_blas_daxpy(-1.0, gradient, state->q);
  
  gsl_blas_ddot(state->p, state->q, &D);
  if(D <= 0.0)
    /* If the curvature condition fails, B (or S) is set to identity. */
    gsl_matrix_set_identity(state->H);
  else
    bfgs_invhess_update(state->p, state->q, state->Sq, state->H);
  
  /* update the current iterate and its function and gradient values */
  gsl_blas_dcopy(state->x_plus, x);
  gsl_blas_dcopy(state->gradient_plus, gradient);
  *f = state->f_plus;
  gsl_blas_dcopy(state->p, dx);
  
  return GSL_SUCCESS;
}

static int bfgs_restart (void *vstate)
{
  bfgs_state_t *state = (bfgs_state_t *)vstate;
  
  gsl_matrix_set_identity(state->H);
  
  return GSL_SUCCESS;
}

static const gsl_multimin_fdfminimizer_type bfgs_hess_f_type = 
{
  "bfgs_hess_f",
  sizeof (bfgs_state_t),
  &bfgs_f_alloc,
  &bfgs_set,
  &bfgs_hess_iterate, 
  &bfgs_restart, 
  &bfgs_free
};

static const gsl_multimin_fdfminimizer_type bfgs_f_type = 
{
  "bfgs_f", 
  sizeof (bfgs_state_t),
  &bfgs_f_alloc,
  &bfgs_set,
  &bfgs_iterate, 
  &bfgs_restart, 
  &bfgs_free
};

static const gsl_multimin_fdfminimizer_type bfgs_hess_mt_type = 
{
  "bfgs_hess_mt",
  sizeof (bfgs_state_t),
  &bfgs_mt_alloc,
  &bfgs_set,
  &bfgs_hess_iterate,
  &bfgs_restart,
  &bfgs_free
};

static const gsl_multimin_fdfminimizer_type bfgs_mt_type = 
{
  "bfgs_mt",
  sizeof (bfgs_state_t),
  &bfgs_mt_alloc,
  &bfgs_set,
  &bfgs_iterate,
  &bfgs_restart,
  &bfgs_free
};

const gsl_multimin_fdfminimizer_type
* gsl_ext_multimin_fdfminimizer_bfgs_hess_f = &bfgs_hess_f_type;

const gsl_multimin_fdfminimizer_type
* gsl_ext_multimin_fdfminimizer_bfgs_f = &bfgs_f_type;

const gsl_multimin_fdfminimizer_type
* gsl_ext_multimin_fdfminimizer_bfgs_hess_mt = &bfgs_hess_mt_type;

const gsl_multimin_fdfminimizer_type
* gsl_ext_multimin_fdfminimizer_bfgs_mt = &bfgs_mt_type;
